#!/usr/bin/python2
# coding: utf8
from indigo import *
from indigo_renderer import *
import os, sys, cStringIO, tempfile
from rdkit import Chem
from fm_parser import exact_mass

class Mol2pic():
    def __init__(self, mol_string = '', width = 600, height = 600):
        self.idg = Indigo()
        self.renderer = IndigoRenderer(self.idg)
        self.idg.setOption("render-comment-position", "top")
        self.idg.setOption("render-image-size", width, height)
#        self.idg.setOption("render-background-color", 1.0, 1.0, 1.0)
        self.idg.setOption('render-label-mode', 'forcehide')
        self.idg.setOption('render-margins', 5, 5)
        self.idg.setOption('render-coloring', True)
        self.idg.setOption('render-stereo-old-style', True)

        if mol_string:
            self.set_mol(mol_string)

    def set_mol(self, string, fmt = 'mol'):
        #print string
        string = string.replace('\\n', '\n')
        if os.path.exists(string):
            self.mol = self.idg.loadMoleculeFromFile(string)
        else:
            self.mol = self.idg.loadMolecule(string)
        self.mol.foldHydrogens()

    def to_pic(self, out_fmt = 'svg', auto_layout = False, filename = None):
        self.idg.setOption("render-output-format", out_fmt)
        if auto_layout:
            self.mol.layout()
        if filename:
            try:
                self.renderer.renderToFile(self.mol, filename)
                return filename
            except Exception, e:
                print e
                pass
        return self.renderer.renderToBuffer(self.mol).tostring()

    def to_mol(self, auto_layout = False):
        if auto_layout:
            self.mol.layout()
        return self.mol.molfile()

    def to_smi(self):
        return self.mol.smiles().split(' |')[0]

    @property
    def formula(self):
        return self.mol.grossFormula().replace(' ', '')

    @property
    def exact_mass(self):
        # take '*' as elememt 'H'
        return exact_mass(self.formula) + self.to_mol().count('*') * 1.007825

def sub_filter(sub, mol):
    ptn = Chem.MolFromSmiles(sub)
    if ptn and mol.HasSubstructMatch(ptn):
        return True
    return False

def filter_sdf(sdf, filters = []):
    #open('a.sdf', 'w').write(sdf)
    results = []
    if not filters:
        num_all = sdf.count('M  END')
        if num_all == 0:
            num_all = sdf.count('$$$$')
        return {'sdf': sdf, 'num_all': num_all, 'num_hits': num_all}
    if os.path.exists(sdf):
        sdffile = sdf
    else:
        # store sdf to temp file
        sdffile = tempfile.mkstemp(text = True)[1]
        open(sdffile, 'w').write(sdf)
    mols = Chem.SDMolSupplier(sdffile)
    hits = []
    for num_all, mol in enumerate(mols):
        for filter in filters:
            if filter(mol):
                hits.append(mol)
    num_all += 1
    sio = cStringIO.StringIO()
    w = Chem.SDWriter(sio)
    for mol in hits:
        w.write(mol)
    w.flush()
    w.close()
    # remove the temp sdffile
    try:
        os.remove(sdffile)
    except:
        pass
    return {'sdf': sio.getvalue(), 'num_all': num_all, 'num_hits': len(hits)}

if __name__ == '__main__':
    sdf = filter_sdf(open('..\\..\\a.sdf').read(), [lambda mol: sub_filter('C1C(S(N)(=O)=O)=CC=C(N)C=1', mol)])
    open('b.sdf', 'w').write(sdf)
    raw_input()
